{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Closed-Loop Evaluation\n",
    "In this notebook you are going to evaluate a CNN-based policy to control the SDV with a protocol named *closed-loop* evaluation.\n",
    "\n",
    "**Note: this notebook assumes you've already run the [training notebook](./train.ipynb) and stored your model successfully.**\n",
    "\n",
    "## What is closed-loop evaluation?\n",
    "In closed-loop evaluation the model is in **full control of the SDV**. At each time step, we predict the future trajectory and then move the AV to the first of the model's predictions. \n",
    "\n",
    "We refer to this process with the terms **forward-simulate** or **unroll**.\n",
    "\n",
    "![closed-loop](../../docs/images/planning/closed-loop.svg)\n",
    "\n",
    "\n",
    "## What can we use closed-loop evaluation for?\n",
    "Closed-loop is crucial to asses a model's capabilities before deploying it in the real world. **Ideally, we would test the model on the road in the real world**. However, this is clearly very expensive and scales poorly. Forward-simulation is an attempt to evaluate the system in a setting which is as close as possible to a real road test on the same route.\n",
    "\n",
    "Differently from open-loop, the model is now in full control of the SDV and predictions are used to compute the future location of the SDV.\n",
    "\n",
    "## Is closed-loop evaluation enough?\n",
    "Closed-loop evaluation is an important step forward towards evaluating how our policy would perform on the road. However, it still has some limitations.\n",
    "\n",
    "The critical one is the **non-reactivity of other traffic participants**. In fact, while the SDV is now controlled by our policy, other actors are still being replayed from the original log. In this setting, a chasing car will not slow down if our policy choses a different speed profile for the SDV, resulting in a rear collision that wouldn't obviously occur in the real world.\n",
    "\n",
    "For this reason, closed-loop evaluation is only accurate for the first few seconds of forward simulation. This can be mitigated when enough data exists for the task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from prettytable import PrettyTable\n",
    "\n",
    "from l5kit.configs import load_config_data\n",
    "from l5kit.data import LocalDataManager, ChunkedDataset\n",
    "from l5kit.dataset import EgoDataset\n",
    "from l5kit.rasterization import build_rasterizer\n",
    "\n",
    "from l5kit.simulation.dataset import SimulationConfig\n",
    "from l5kit.simulation.unroll import ClosedLoopSimulator\n",
    "from l5kit.cle.closed_loop_evaluator import ClosedLoopEvaluator, EvaluationPlan\n",
    "from l5kit.cle.metrics import (CollisionFrontMetric, CollisionRearMetric, CollisionSideMetric,\n",
    "                               DisplacementErrorL2Metric, DistanceToRefTrajectoryMetric)\n",
    "from l5kit.cle.validators import RangeValidator, ValidationCountingAggregator\n",
    "\n",
    "from l5kit.visualization.visualizer.zarr_utils import simulation_out_to_visualizer_scene\n",
    "from l5kit.visualization.visualizer.visualizer import visualize\n",
    "from bokeh.io import output_notebook, show\n",
    "from l5kit.data import MapAPI\n",
    "\n",
    "from collections import defaultdict\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare data path and load cfg\n",
    "\n",
    "By setting the `L5KIT_DATA_FOLDER` variable, we can point the script to the folder where the data lies.\n",
    "\n",
    "Then, we load our config file with relative paths and other configurations (rasteriser, training params...)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set env variable for data\n",
    "os.environ[\"L5KIT_DATA_FOLDER\"] = \"/tmp/l5kit_data\"\n",
    "dm = LocalDataManager(None)\n",
    "# get config\n",
    "cfg = load_config_data(\"./config.yaml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = \"/tmp/planning_model.pt\"\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = torch.load(model_path).to(device)\n",
    "model = model.eval()\n",
    "torch.set_grad_enabled(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the evaluation data\n",
    "Differently from training and open loop evaluation, this setting is intrinsically sequential. As such, we won't be using any of PyTorch's parallelisation functionalities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ===== INIT DATASET\n",
    "eval_cfg = cfg[\"val_data_loader\"]\n",
    "rasterizer = build_rasterizer(cfg, dm)\n",
    "eval_zarr = ChunkedDataset(dm.require(eval_cfg[\"key\"])).open()\n",
    "eval_dataset = EgoDataset(cfg, eval_zarr, rasterizer)\n",
    "print(eval_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define some simulation properties\n",
    "We define here some common simulation properties such as the length of the simulation and how many scene to simulate.\n",
    "\n",
    "**NOTE: these properties have a significant impact on the execution time. We suggest you to increase them only if your setup includes a GPU**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_scenes_to_unroll = 10\n",
    "num_simulation_steps = 50"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Closed-loop simulation\n",
    "\n",
    "We define a closed-loop simulation that drives the SDV for `num_simulation_steps` steps while using the log-replay agents.\n",
    "\n",
    "Then, we unroll the selected scenes.\n",
    "The simulation output contains all the information related to the scene, including the annotated and simulated positions, states, and trajectories of the SDV and the agents.  \n",
    "If you want to know more about what the simulation output contains, please refer to the source code of the class `SimulationOutput`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==== DEFINE CLOSED-LOOP SIMULATION\n",
    "sim_cfg = SimulationConfig(use_ego_gt=False, use_agents_gt=True, disable_new_agents=True,\n",
    "                           distance_th_far=500, distance_th_close=50, num_simulation_steps=num_simulation_steps,\n",
    "                           start_frame_index=0, show_info=True)\n",
    "\n",
    "sim_loop = ClosedLoopSimulator(sim_cfg, eval_dataset, device, model_ego=model, model_agents=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==== UNROLL\n",
    "scenes_to_unroll = list(range(0, len(eval_zarr.scenes), len(eval_zarr.scenes)//num_scenes_to_unroll))\n",
    "sim_outs = sim_loop.unroll(scenes_to_unroll)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Closed-loop metrics\n",
    "In this setting **metrics are particularly challenging**.\n",
    "In fact, we would like to penalise some of the simulation drift (e.g. going off road or in the opposite lane) while at the same time allow others (e.g. different speed profiles).\n",
    "\n",
    "\n",
    "### Collisions\n",
    "Our SDV should avoid collisions with other agents. In this example, we won't distinguish between collisions caused by non-reactivity of other agents and actual collisions, and we will simply report them all categorised by where they occurred (front, rear and side with respect to the AV).\n",
    "\n",
    "However, if we only considered collisions, our SDV could pass all tests by simply driving off-road or in a different lane.\n",
    "\n",
    "\n",
    "### Distance from reference trajectory\n",
    "To address the issue presented above, we require our SDV to loosely stick to the original trajectory in the data. By setting the right threshold on the distance we can allow for different speed profiles and small steerings, while pensalising large deviations like driving off-road.\n",
    "\n",
    "We can do so by computing the distance between the simulated trajectory and the annotated one in world coordinates.\n",
    "\n",
    "\n",
    "### Displacement error (L2)\n",
    "In addition, we can measure the displacement between each point of the simulated trajectory and the corresponding annotated one in world coordinates.\n",
    "\n",
    "With this metric, we can spot large deviations between the speed of the simulated policy and the annotated one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = [DisplacementErrorL2Metric(),\n",
    "           DistanceToRefTrajectoryMetric(),\n",
    "           CollisionFrontMetric(),\n",
    "           CollisionRearMetric(),\n",
    "           CollisionSideMetric()]\n",
    "\n",
    "validators = [RangeValidator(\"displacement_error_l2\", DisplacementErrorL2Metric, max_value=30),\n",
    "              RangeValidator(\"distance_ref_trajectory\", DistanceToRefTrajectoryMetric, max_value=4),\n",
    "              RangeValidator(\"collision_front\", CollisionFrontMetric, max_value=0),\n",
    "              RangeValidator(\"collision_rear\", CollisionRearMetric, max_value=0),\n",
    "              RangeValidator(\"collision_side\", CollisionSideMetric, max_value=0)]\n",
    "\n",
    "intervention_validators = [\"displacement_error_l2\",\n",
    "                           \"distance_ref_trajectory\",\n",
    "                           \"collision_front\",\n",
    "                           \"collision_rear\",\n",
    "                           \"collision_side\"]\n",
    "\n",
    "cle_evaluator = ClosedLoopEvaluator(EvaluationPlan(metrics=metrics,\n",
    "                                                   validators=validators,\n",
    "                                                   composite_metrics=[],\n",
    "                                                   intervention_validators=intervention_validators))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantitative evaluation\n",
    "\n",
    "We can now compute the metric evaluation, collect the results and aggregate them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cle_evaluator.evaluate(sim_outs)\n",
    "validation_results = cle_evaluator.validation_results()\n",
    "agg = ValidationCountingAggregator().aggregate(validation_results)\n",
    "cle_evaluator.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reporting errors from the closed-loop\n",
    "\n",
    "We can now report the metrics and plot them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fields = [\"metric\", \"value\"]\n",
    "table = PrettyTable(field_names=fields)\n",
    "\n",
    "values = []\n",
    "names = []\n",
    "\n",
    "for metric_name in agg:\n",
    "    table.add_row([metric_name, agg[metric_name].item()])\n",
    "    values.append(agg[metric_name].item())\n",
    "    names.append(metric_name)\n",
    "\n",
    "print(table)\n",
    "\n",
    "plt.bar(np.arange(len(names)), values)\n",
    "plt.xticks(np.arange(len(names)), names, rotation=60, ha='right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Qualitative evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualise the closed-loop\n",
    "\n",
    "We can visualise the scenes we have obtained previously. \n",
    "\n",
    "**The policy is now in full control of the SDV as this moves through the annotated scene.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "output_notebook()\n",
    "mapAPI = MapAPI.from_cfg(dm, cfg)\n",
    "for sim_out in sim_outs: # for each scene\n",
    "    vis_in = simulation_out_to_visualizer_scene(sim_out, mapAPI)\n",
    "    show(visualize(sim_out.scene_id, vis_in))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pre-trained model results\n",
    "\n",
    "We include here the unroll of one scene using one of our pre-trained model. The controlled SDV can stick to the correct lane and stops successfully for a red light. \n",
    "\n",
    "Comparing this result with the one from the [open-loop notebook](./open_loop_test.ipynb) we can notice some differences in the speed profile chosen by the model.\n",
    "\n",
    "\n",
    "![SegmentLocal](../../docs/images/planning/out_9_closed.gif \"segment\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.9"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
